import os
from typing import Tuple
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import gridspec
from matplotlib import cm

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
plt.rcParams['figure.max_open_warning'] = 50

WORKSPACE = os.path.dirname(__file__)
AROUND = True
COUNTER = 0

def model1(sample: np.ndarray) -> np.ndarray:
    result = np.mean(sample, axis=-1) * 2 - 1
    if AROUND: result = np.around(result)
    result = np.array([result])
    return np.reshape(result, (result.shape[-1], 1))

def model2(sample: np.ndarray) -> np.ndarray:
    result = np.median(sample, axis=-1) * 2 - 1
    if AROUND: result = np.around(result)
    result = np.array([result])
    return np.reshape(result, (result.shape[-1], 1))

def model3(sample: np.ndarray) -> np.ndarray:
    result = np.max(sample, axis=-1) + np.min(sample, axis=-1) - 1
    if AROUND: result = np.around(result)
    result = np.array([result])
    return np.reshape(result, (result.shape[-1], 1))

def model4(sample: np.ndarray) -> np.ndarray:
    n = sample.shape[-1]
    result = (1 + 1 / n) * np.max(sample, axis=-1) - 1 / n
    if AROUND: result = np.around(result)
    result = np.array([result])
    return np.reshape(result, (result.shape[-1], 1))

def model5(sample: np.ndarray) -> np.ndarray:
    n = sample.shape[-1]
    result = (1 + 1 / (2 * n - 1)) * (np.max(sample, axis=-1) - 1 / (2 * n))
    if AROUND: result = np.around(result)
    result = np.array([result])
    return np.reshape(result, (result.shape[-1], 1))

def calculate_sample(sample: np.ndarray) -> np.matrix:
    result = np.hstack((model1(sample), model2(sample), model3(sample), model4(sample), model5(sample)))
    return np.matrix(result)

def random_samples(x: int=1000, n: int=10, m: int=200) -> Tuple[np.matrix, np.matrix]:
    global COUNTER
    samples = np.random.choice(x, n) + 1
    for _ in range(m - 1):
        samples = np.vstack((samples, np.random.choice(x, n) + 1))
    dirpath = os.path.join(WORKSPACE, "sample_x={:d}_n={:d}_m={:d}_{:d}".format(x, n, m, COUNTER))
    COUNTER += 1
    os.makedirs(dirpath, exist_ok=True)
    np.savetxt(os.path.join(dirpath, "sample.csv"), samples, delimiter=',', fmt='%04d')
    results = calculate_sample(samples)
    results_mean = np.mean(results, axis=0)
    results_mean_difference = results_mean - x
    results_std = np.std(results, axis=0)
    samples_max = np.max(samples, axis=-1)
    samples_max = np.tile(samples_max, (5,1)).T
    results_difference = results - samples_max
    results_difference = np.where(results_difference < 0, 1, 0)
    results_difference = np.sum(results_difference, axis=0)
    np.savetxt(os.path.join(dirpath, 'result.csv'), np.vstack((results, results_mean, results_mean_difference, results_std, results_difference)), delimiter=',', fmt='%.1f')
    
    fig = plt.figure(tight_layout=True)
    gs = gridspec.GridSpec(2, 6)
    
    ax1 = fig.add_subplot(gs[0, 0:2])
    ax1.hist(results[:, 0], bins=10)
    ax1.set_xlabel("模型1")
    
    ax2 = fig.add_subplot(gs[0, 2:4])
    ax2.hist(results[:, 1], bins=10)
    ax2.set_xlabel("模型2")
    
    ax3 = fig.add_subplot(gs[0, 4:6])
    ax3.hist(results[:, 2], bins=10)
    ax3.set_xlabel("模型3")
    
    ax4 = fig.add_subplot(gs[1, 1:3])
    ax4.hist(results[:, 3], bins=10)
    ax4.set_xlabel("模型4")
    
    ax5 = fig.add_subplot(gs[1, 3:5])
    ax5.hist(results[:, 4], bins=10)
    ax5.set_xlabel("模型5")
    
    fig.align_labels()
    plt.savefig(os.path.join(dirpath, 'result.png'))
    
    return results_mean_difference, results_std

def task1():
    path1 = os.path.join(WORKSPACE, 'sample1', "sample.csv")
    path2 = os.path.join(WORKSPACE, 'sample2', "sample.csv")
    sample1 = np.loadtxt(path1, delimiter=',', dtype=float)
    sample2 = np.loadtxt(path2, delimiter=',', dtype=float)
    result1 = calculate_sample(sample1)
    result2 = calculate_sample(sample2)
    print("ptp1: ", np.ptp(result1), "\nptp2: ", np.ptp(result2), "\nptp: ", np.abs(result1 - result2))
    np.savetxt(os.path.join(os.path.dirname(path1), 'result.csv'), result1, delimiter=',', fmt='%04d')
    np.savetxt(os.path.join(os.path.dirname(path2), 'result.csv'), result2, delimiter=',', fmt='%04d')

if __name__ == "__main__":
    task1()
    n_tuple = (500, 100, 50, 10)
    m_tuple = (3200, 1600, 800, 400, 200)
    results_mean_difference_all = np.zeros((len(m_tuple), len(n_tuple), 5))
    results_std_all = np.zeros((len(m_tuple), len(n_tuple), 5))
    random_samples()
    for m_id, m in enumerate(m_tuple):
        for n_id, n in enumerate(n_tuple):
            print(m_id, n_id)
            results_mean_difference, results_std = random_samples(n=n, m=m)
            results_mean_difference = np.array(results_mean_difference)
            results_std = np.array(results_std)
            for model_id in range(5):
                results_mean_difference_all[m_id][n_id][model_id] += abs(results_mean_difference[:, model_id])
                results_std_all[m_id][n_id][model_id] += results_std[:, model_id]
    N, M = np.meshgrid(n_tuple, m_tuple)
    for model_id in range(5):
        fig = plt.figure(tight_layout=True, figsize=(16, 8))
        gs = gridspec.GridSpec(4, 2)
        
        ax1 = fig.add_subplot(gs[0:2, 1], projection='3d')
        surf1 = ax1.plot_surface(N, M, results_mean_difference_all[:, :, model_id], cmap=cm.coolwarm)
        ax1.set_title("模型{:d}平均值的误差随样本大小n、样本数量m变化".format(model_id + 1))
        ax1.set_xlabel("n")
        ax1.set_ylabel("m")
        fig.colorbar(surf1, shrink=0.5, aspect=5)
        
        ax2 = fig.add_subplot(gs[2:4, 1], projection='3d')
        surf2 = ax2.plot_surface(N, M, results_std_all[:, :, model_id], cmap=cm.coolwarm)
        ax2.set_title("模型{:d}标准差随样本大小n、样本数量m变化".format(model_id + 1))
        ax2.set_xlabel("n")
        ax2.set_ylabel("m")
        fig.colorbar(surf2, shrink=0.5, aspect=5)
        
        for row_id in range(4):
            ax = fig.add_subplot(gs[row_id, 0])
            results = results_mean_difference_all if row_id < 2 else results_std_all
            str1 = "平均值的误差" if row_id < 2 else "标准差"
            if row_id % 2 == 0:
                for m_id, m in enumerate(m_tuple):
                    ax.plot(n_tuple, results[m_id, :, model_id], label="$m={}$".format(m))
                str2 = "随样本大小n变化"
            else:
                for n_id, n in enumerate(n_tuple):
                    ax.plot(m_tuple, results[:, n_id, model_id], label="$n={}$".format(n))
                str2 = "随样本数量m变化"
            ax.set_title("模型{:d}".format(model_id+1) + str1 + str2)
            
            plt.legend()
        
        plt.savefig(os.path.join(WORKSPACE, 'result_{:d}.png'.format(model_id)))